"""
SSH连接管理模块
"""

import subprocess
import threading
import time
import os
import signal
from typing import Optional, Callable
from config import SSH_PORT, SSH_USERNAME, HEARTBEAT_INTERVAL

class SSHManager:
    def __init__(self):
        self.tunnel_process: Optional[subprocess.Popen] = None
        self.heartbeat_thread: Optional[threading.Thread] = None
        self.is_running = False
        self.heartbeat_callback: Optional[Callable] = None
    
    def create_reverse_tunnel(self, relay_host: str, relay_port: int, tunnel_port: int, 
                            ssh_user: str = SSH_USERNAME, local_port: int = SSH_PORT) -> bool:
        """
        创建SSH反向隧道
        
        Args:
            relay_host: 中继服务器地址
            relay_port: 中继服务器SSH端口（通常是22）
            tunnel_port: 隧道端口
            ssh_user: SSH用户名
            local_port: 本地SSH端口
            
        Returns:
            是否成功创建隧道
        """
        try:
            # 构建SSH反向隧道命令
            # ssh -R tunnel_port:localhost:local_port -N -o StrictHostKeyChecking=no user@relay_host
            cmd = [
                "ssh",
                "-R", f"{tunnel_port}:localhost:{local_port}",
                "-N",  # 不执行远程命令
                "-o", "StrictHostKeyChecking=no",  # 不检查主机密钥
                "-o", "UserKnownHostsFile=/dev/null",  # 不保存主机密钥
                "-o", "ServerAliveInterval=30",  # 30秒发送一次保活消息
                "-o", "ServerAliveCountMax=3",   # 最多3次保活失败
                f"{ssh_user}@{relay_host}"
            ]
            
            # 如果指定了SSH端口且不是默认的22
            if relay_port != 22:
                cmd.extend(["-p", str(relay_port)])
            
            print(f"执行命令: {' '.join(cmd)}")
            
            # 启动SSH隧道进程
            self.tunnel_process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                preexec_fn=os.setsid  # 创建新的进程组
            )
            
            # 等待一小段时间检查进程是否正常启动
            time.sleep(2)
            
            if self.tunnel_process.poll() is None:
                print(f"SSH反向隧道已建立: {relay_host}:{tunnel_port} -> localhost:{local_port}")
                return True
            else:
                stdout, stderr = self.tunnel_process.communicate()
                print(f"SSH隧道启动失败: {stderr.decode()}")
                return False
                
        except Exception as e:
            print(f"创建SSH隧道时发生错误: {str(e)}")
            return False
    
    def start_heartbeat(self, callback: Callable):
        """
        启动心跳线程
        
        Args:
            callback: 心跳回调函数
        """
        self.heartbeat_callback = callback
        self.is_running = True
        self.heartbeat_thread = threading.Thread(target=self._heartbeat_worker, daemon=True)
        self.heartbeat_thread.start()
    
    def _heartbeat_worker(self):
        """心跳工作线程"""
        while self.is_running:
            try:
                if self.heartbeat_callback:
                    result = self.heartbeat_callback()
                    if result.get("status") != "alive":
                        print(f"心跳失败: {result.get('message', '未知错误')}")
                        break
                
                time.sleep(HEARTBEAT_INTERVAL)
            except Exception as e:
                print(f"心跳线程错误: {str(e)}")
                break
    
    def stop_tunnel(self):
        """停止SSH隧道"""
        self.is_running = False
        
        if self.tunnel_process:
            try:
                # 终止进程组
                os.killpg(os.getpgid(self.tunnel_process.pid), signal.SIGTERM)
                
                # 等待进程结束
                try:
                    self.tunnel_process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    # 如果5秒内没有结束，强制杀死
                    os.killpg(os.getpgid(self.tunnel_process.pid), signal.SIGKILL)
                    self.tunnel_process.wait()
                
                print("SSH隧道已停止")
            except Exception as e:
                print(f"停止SSH隧道时发生错误: {str(e)}")
            finally:
                self.tunnel_process = None
        
        if self.heartbeat_thread and self.heartbeat_thread.is_alive():
            self.heartbeat_thread.join(timeout=2)
    
    def is_tunnel_active(self) -> bool:
        """检查隧道是否活跃"""
        return self.tunnel_process is not None and self.tunnel_process.poll() is None
    
    def connect_to_target(self, relay_host: str, tunnel_port: int, 
                         ssh_user: str = SSH_USERNAME) -> bool:
        """
        连接到目标主机
        
        Args:
            relay_host: 中继服务器地址
            tunnel_port: 隧道端口
            ssh_user: SSH用户名
            
        Returns:
            是否成功连接
        """
        try:
            # 构建SSH连接命令
            cmd = [
                "ssh",
                "-o", "StrictHostKeyChecking=no",
                "-o", "UserKnownHostsFile=/dev/null",
                "-p", str(tunnel_port),
                f"{ssh_user}@{relay_host}"
            ]
            
            print(f"连接到目标主机: {' '.join(cmd)}")
            
            # 使用os.system直接执行，这样可以保持交互式终端
            result = os.system(' '.join(cmd))
            return result == 0
            
        except Exception as e:
            print(f"连接目标主机时发生错误: {str(e)}")
            return False

